#将split_train_val_yolo.py生成的train.txt val.txt两个文件进行读取，
#并查询对应json文件的backgroundNumber信息
#并读取根目录下的backgrounds图片文件, 进行对应的差分

import argparse
from pathlib import Path
import xml.etree.ElementTree as ET
import cv2
import os
from tqdm import tqdm

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="generateBackgroundTxt")
    parser.add_argument('--root-dir', default="", type=str)
    return parser.parse_args()

def main():
    args = Parse_Arguments()
    root_dir = args.root_dir
    background_imgs_path = args.root_dir + "/backgrounds/" #背景图存放路径
    background_residual_path = args.root_dir + "/background_residual/" #背景差分图存放路径
    os.makedirs(background_residual_path, exist_ok=True)
    images_path = []
    with open(root_dir + '/dataSet_path/train.txt', 'r') as f:
        images_path += f.read().strip().splitlines()
    with open(root_dir + '/dataSet_path/val.txt', 'r') as f:
        images_path += f.read().strip().splitlines()

    for image_path in tqdm(images_path, desc="正在飞速的生成背景差分(●'◡'●)"):
        xml_filename = Path(image_path).with_suffix('.xml').name
        xml_file_path = root_dir + f"/Annotations/{xml_filename}"
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        try:
            value = root.find('backgroundNumber').text
            background_img_path = background_imgs_path + value + ".jpg"
            image = cv2.imread(image_path)
            background_img = cv2.imread(background_img_path)
            diff_image = cv2.absdiff(background_img, image)
            diff_image = cv2.cvtColor(
                diff_image, cv2.COLOR_BGR2GRAY)
            store_residual_path = background_residual_path + str(Path(image_path).stem) + "_residual.jpg"
            cv2.imwrite(store_residual_path, diff_image)
        except Exception as e:
            print(f"错误详情: {e}")
            print("跳过该图片")
            print(f"请检查{xml_filename}和对应图片, 有可能是该文件里没有加入backgroundNumber标签")


if __name__ == "__main__":
    main()
